import numpy as np
from ptflops import get_model_complexity_info
import torch
import torch.nn as nn
import torch.optim as optim
from utils.config import *
from utils.init_weights import init_weights
from core.gc_net import GCNET
from dataloader.data_loader import train_loader, validate_loader
from train import train
from validate import validate
from save import save
from torch.utils.tensorboard import SummaryWriter

def main():
    # 创建tensorboard log文件
    writer = SummaryWriter(log_dir='logs')

    # 打印数据集名称
    print('dataset:', dataset)

    # 创建模型实例
    model = GCNET(height, width, max_disp)
    model = nn.DataParallel(model, device_ids=[0, 1]).to(device)
    torch.backends.cudnn.benchmark = True

    # 打印网络参数量
    print("Number of parameters: ", sum(p.numel() for p in model.parameters()))
    if dataset == 'usvinland':
        # 使用ptflops模块计算复杂度
        prepare_input = lambda _: {"img_left": torch.FloatTensor(1, 3, 320, 640).to('cuda'), "img_right": torch.FloatTensor(1, 3, 320, 640).to('cuda')}
        macs, params = get_model_complexity_info(model.module, input_res=(3, 320, 640), input_constructor=prepare_input, print_per_layer_stat=False, verbose=False)
        print(f'ptflops: {{ macs: {macs}, params: {params} }}')

    # 打印网络结构
    # print(model)

    # 权重初始化
    if load_pretrained:
        pretrained_path = './pretrained/sceneflow.pth'
        state = torch.load(pretrained_path)
        model.load_state_dict(state['state_dict'])
    else:
        init_weights(model)

    # 损失函数
    criterion = torch.nn.L1Loss().to(device)

    # 优化器
    optimizer = optim.RMSprop(model.parameters(), lr=0.001)

    # 训练与验证
    step = 0
    best_error = np.inf
    model.cuda()
    for epoch in range(NUM_EPOCHS):
        # 训练
        model.train()
        step = train(model, train_loader, criterion, optimizer, step, writer)

        # 验证
        model.eval()
        error = validate(model, validate_loader, epoch, writer)

        # 保存最优模型
        best_error = save(model, optimizer, epoch, step, error, best_error)

    writer.close()

if __name__ == '__main__':
    main()
